require 'random'
require 'cipher'
require 'binstr'

module AllOrNothing

  class Packager
    KNOWN_KEY = 0x69.chr     # from PyCrypto source
    TEXT_PAD_CHAR = ' '
    MIN_PACKAGE_BLOCKS = 3   # 1+ for data, 1 for pad_length, 1 for session key

    def self.known_key(key_length)
      KNOWN_KEY * key_length
    end

    def initialize(cipher, key_length = 16)
      @cipher = cipher
      @key_length = key_length
    end

    def known_key
      self.class.known_key(@key_length)
    end

    def pack(text) # returns package (array of strings)
      pad_length = @cipher.block_size - 
        ((text.length - 1 % @cipher.block_size) + 1)

      text = text + TEXT_PAD_CHAR * pad_length

      puts "text = \"#{text}\""

      random_key = Random.pseudo_bytes(@key_length)

      puts "random_key = #{BinStr.bytes random_key}"

      blocks = []
      hashes = []

      # blocks count from 1 (TODO: why not 0?)
      i = 1

      while !text.empty?
        block = text.slice!(0, @cipher.block_size)
        puts "block = \"#{block}\""
        warn "block size mismatch" if block.size != @cipher.block_size

        # encrypt i with random_key
        cipherblock = @cipher.encrypt(BinStr.from_int(i, @cipher.block_size), 
                                      random_key)
        # block XOR encrypted i
        tmp = BinStr.to_int(block, @cipher.block_size) ^ 
          BinStr.to_int(cipherblock, @cipher.block_size)
        blocks << tmp

        # encrypt (above result XOR plain i) with known_key
        hashes << @cipher.encrypt(BinStr.from_int(tmp ^ i, @cipher.block_size),
                                  self.known_key)
        i += 1
      end

      i += 1
      cipherblock = @cipher.encrypt(BinStr.from_int(i, @cipher.block_size),
                                    random_key)
      tmp = pad_length ^ BinStr.to_int(cipherblock, @cipher.block_size)
      blocks << tmp
      hashes << @cipher.encrypt(BinStr.from_int(tmp ^ i, @cipher.block_size),
                                self.known_key)

      puts "blocks.length = #{blocks.length}"

      xor_reduction = hashes.inject(0) { |memo, str|
        memo ^ BinStr.to_int(str, @cipher.block_size)
      }

      blocks << (BinStr.to_int(random_key, @cipher.block_size) ^ xor_reduction)
      blocks.map { |int| BinStr.from_int(int, @cipher.block_size) }
    end

    def unpack(package) # returns text
      if package.length < MIN_PACKAGE_BLOCKS
        raise "package.length must be at least #{MIN_PACKAGE_BLOCKS}"
      end
      package = package.map { |str| BinStr.to_int(str, @cipher.block_size) }
      
      hashes = []
      package.each_with_index { |int, index|
        tmp = BinStr.from_int(int ^ (index + 1), @cipher.block_size)
        hashes << BinStr.to_int(@cipher.encrypt(tmp, self.known_key),
                                @cipher.block_size)
      }
      derived_key = package.last ^ package.inject(0) { |memo, int| memo ^ int }
      derived_key = BinStr.from_int(derived_key, @key_length)

      puts "derived_key =\n#{BinStr.bytes derived_key}"

      parts = []
      package.each_with_index { |int, index|
        cipherblock = @cipher.encrypt(BinStr.from_int(int, @cipher.block_size),
                                      derived_key)
        parts << (int ^ BinStr.to_int(cipherblock, @cipher.block_size))
      }

      pad_length = parts.pop.to_i
      puts pad_length
      parts.map { |int| 
        BinStr.from_int(int, @cipher.block_size) 
      }.join.chop_rear(pad_length)
    end
  end
end
